package statistics;

import java.util.ArrayList;
import java.util.List;

import dnaModels.DNASequence;

import population.Individual;

/**
 * A crazy tree sampler for experimental use only. Creates a select panel of individuals and filters for only those haplotypes to
 * select for contemporary samples. 
 * @author brendan
 *
 */
public class PanelTreeSampler extends SerialTreeSampler {

	public PanelTreeSampler(int[] samplingTimes, int[] sampleSizes) {
		super(samplingTimes, sampleSizes);
	}

	
	public void collect(Collectible pop) {
		//System.out.println("Serial tree sampler is collecting");
		
		if (pop.getCurrentGenNumber()==lastGenCollected) {
			//System.out.println("Returning since last gen collected = " + pop.getCurrentGenNumber());
			return;
		}
		
		if (firstGenCalled<0)
			firstGenCalled = pop.getCurrentGenNumber();
		
		//System.out.println("Not returning. on gen " + pop.getCurrentGenNumber());
		lastGenCollected = pop.getCurrentGenNumber();
		
		if ((firstGenCalled-collectionOffset)==pop.getCurrentGenNumber() || (pop.getCurrentGenNumber()-(firstGenCalled-collectionOffset))%collectionFrequency==0) {
			if (collecting) {
				System.out.println("We're still collecting but we're trying to initiate a new collection period");
				return;
			}
			System.out.println("\n\n ** Gen : " + pop.getCurrentGenNumber() + " starting new serial data collection, N=" + pop.size());
			tips = new ArrayList<Individual>();
			startGen = pop.getCurrentGenNumber();
			nextSampleIndex = 0;
			collecting = true;
		} 
		
		
		if (collecting) { 
			int elapsedGens = pop.getCurrentGenNumber()-startGen;
			
			if ( samplingTimes[nextSampleIndex]==elapsedGens) { 
				
				if (nextSampleIndex == (samplingTimes.length-1) ) {
					int matchingRegionSize = 300;
					//This is the last sampling period so use the crazy new scheme
					
					//Create a panel
					List<Individual> panelInds = createPanel(20, pop, 0, matchingRegionSize);
					
					double proportionFromPanel = 0.20;
					
					//Select sampling individuals whose hap matches the panel
					List<Individual> sample = new ArrayList<Individual>();
					while (sample.size() < (int)Math.round((double)sampleSizes[nextSampleIndex]*proportionFromPanel)) {
						Individual prospy = pop.getSample(1).get(0);
						if (! sample.contains(prospy)) {
							if ( hapMatches(prospy.getPrimaryDNA(), panelInds, 0, matchingRegionSize)) {
								sample.add(prospy);
								tips.add(prospy);
								prospy.setPreserve(true);
							}
						}
					}
					
					//Now add more individuals from population at random until sample size is the right size
					while(sample.size() < sampleSizes[nextSampleIndex]) {
						Individual prospy = pop.getSample(1).get(0);
						if (! sample.contains(prospy)) {
							sample.add(prospy);
							tips.add(prospy);
							prospy.setPreserve(true);
						}
					}
					
					
				}
				else {
					List<Individual> sample = pop.getSample(sampleSizes[nextSampleIndex]);
					System.out.println("Gen : " + pop.getCurrentGenNumber() + "  nextSampleIndex: " + nextSampleIndex + " ...  Flagging " + sample.size() + " new individuals for pres");

					for(Individual ind : sample) {
						ind.setPreserve(true);
						tips.add(ind); 
					}
				}
				nextSampleIndex++;
			}

			
			if (nextSampleIndex==samplingTimes.length ) {
				constructTreeFromTips();
				pop.releasePreservedInds();
				collecting = false;
				indsInTree.clear();
				tips.clear();	//It's probably important that we don't retain references to any individuals in the persistent pop
				addDepthLabelsToTree();
				System.out.println("Serial tree sampler has finished another tree, N=" + pop.size());
				fireNewTreeEvent(lastTree);
				nextSampleIndex = 0;
			}
			
		}	
		
	}

	/**
	 * Returns true if the given ind has a DNA sequence that matches any panel member in the given region
	 * @param ind
	 * @param panel
	 * @param regStart
	 * @param regEnd
	 * @return
	 */
	private static boolean hapMatches(DNASequence ind, List<Individual> panel, int regStart, int regEnd) {
		for(Individual panelMember : panel) {
			if (regionMatches(ind, panelMember.getPrimaryDNA(), regStart, regEnd)) {
				return true;
			}
		}
		return false;
	}

	/**
	 * Create a list of X inds all of whom have differing haplotypes in the given region 
	 * @param pop
	 * @return
	 */
	private List<Individual> createPanel(int panelSize, Collectible pop, int regStart, int regEnd) {
		List<Individual> panel = new ArrayList<Individual>();
		panel.add(pop.getSample(1).get(0));
		int count = 0; //Abort after this many tries, since we may not always be able to create a sample of the given size. 
		while (panel.size() < panelSize) {
			Individual ind = pop.getSample(1).get(0);
			if (! panel.contains(ind)) {
				//Search existing panel for any individuals whose dna region is the same as the prospective individual, if none are
				//found add the prospy
				boolean hapContained = false;
				for(int i=0; i<panel.size(); i++) {
					if ( regionMatches(panel.get(i).getPrimaryDNA(), ind.getPrimaryDNA(), regStart, regEnd)) {
						hapContained = true;
						break;
					}
				}
				
				if (! hapContained) {
					panel.add(ind);
					//System.out.println(ind.getPrimaryDNA());
				}
				
				count++;
				if (count > 1000000) {
					throw new IllegalStateException("Cant find enough unique individuals to create panel");
				}
			}
		}
		return panel;
	}
	
	private static boolean regionMatches(DNASequence a, DNASequence b, int start, int end) {
		for(int site = start; site<end; site++) {
			if (a.getBaseChar(site)!=b.getBaseChar(site)) {
				return false;
			}
		}
		return true;
	}
}
